#include <stdio.h>
#include <string.h>
#include <time.h>

#include "crypto/sm2.h"
#include "crypto/sm3.h"

#define IO_BUF_SIZE 256
#define SM2_KEY_HEX_LEN 64

// --------------------------------------------------------------------------------
// Result Storage
// --------------------------------------------------------------------------------
#define USERID_BUFFER_SIZE  64
#define ZSTRING_BUFFER_SIZE 256
#define MESSAGE_BUFFER_SIZE 2048
#define ASN1_DER_BYTE_LEN   70

static uint8_t sm2_storage_uid[USERID_BUFFER_SIZE];
static uint8_t sm2_storage_zstr[ZSTRING_BUFFER_SIZE];
static uint8_t sm2_storage_za[32]; // 256bits
static uint8_t sm2_storage_ma[32];
static uint8_t sm2_storage_message[MESSAGE_BUFFER_SIZE];
static uint8_t sm2_storage_encrypt[MESSAGE_BUFFER_SIZE];
static uint8_t sm2_storage_der[ASN1_DER_BYTE_LEN];
static uint8_t sm2_storage_r[32];
static uint8_t sm2_storage_s[32];

// input buffer
static char input_buf_test_name[IO_BUF_SIZE];
static char input_buf_msg[MESSAGE_BUFFER_SIZE + 1];
static char input_buf_msg_type[10];
static char input_buf_uid[USERID_BUFFER_SIZE];
static char input_buf_r[65];
static char input_buf_s[65];
static char input_buf_digest[65];

/**
 * SM3 Hash tests
 */
void test_sm3_hash_1() {
    sm2_bitstring bstr = sm2_bitstring_wrap(sm2_storage_message, MESSAGE_BUFFER_SIZE, 0);
    sm2_bitstring_set_from_cstr(&bstr, "abc");
    printf("\n========== SM3 Hash Test Case 1 ==========\n");
    uint8_t hash[32];
    sm3_hash(hash, bstr.data, bstr.bit_len);
    printf("* Input hex\n");
    sm2_show_bitstring(&bstr);
    printf("* Hash hex\n");
    sm2_show_bytes(hash, 32, true);
    printf("* Expected hex\n");
    printf("66C7F0F4 62EEEDD9 D1F2D46B DC10E4E2 4167C487 5CF2F7A2 297DA02B 8F4BA8E0\n");
}
void test_sm3_hash_2() {
    sm2_bitstring bstr = sm2_bitstring_wrap(sm2_storage_message, MESSAGE_BUFFER_SIZE, 0);
    sm2_bitstring_set_from_bytes(&bstr, (uint8_t[]){
        0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64,
        0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64,
        0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64,
        0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64,
        0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64,
        0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64,
        0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64,
        0x61, 0x62, 0x63, 0x64, 0x61, 0x62, 0x63, 0x64
    }, 64);
    printf("\n========== SM3 Hash Test Case 2 ==========\n");
    uint8_t hash[32];
    sm3_hash(hash, bstr.data, bstr.bit_len);
    printf("* Input hex\n");
    sm2_show_bitstring(&bstr);
    printf("* Hash hex\n");
    sm2_show_bytes(hash, 32, true);
    printf("* Expected hex\n");
    printf("DEBE9FF9 2275B8A1 38604889 C18E5A4D 6FDB70E5 387E5765 293DCBA3 9C0C5732\n");
}

/**
 * Trim the last return char of a str.
 * @param str get rid of the return chars read by fgets,
 */
void line_trim(char *str) {
    for (int i = (int) strlen(str) - 1; i >= 0; i--) {
        if (str[i] == '\n' || str[i] == '\r') {
            str[i] = '\0';
            break;
        }
    }
}

void read_test_file(const char* input_file, sm2_ctx* ctx) {
    FILE* input = fopen(input_file, "r");
    if (input == NULL) {
        printf("Failed to open file to read %s\n", input_file);
        return;
    }
    const int buf_size = IO_BUF_SIZE;
    char line_buf[IO_BUF_SIZE];
    char op[IO_BUF_SIZE];
    char data_buf [IO_BUF_SIZE];
    while (fgets(line_buf, buf_size, input) != NULL) {
        sscanf(line_buf, "%s %s", op, data_buf);
        if (strcmp(op, "SIGN") == 0 || strcmp(op, "ENCRYPT") == 0) {
            printf("* Running %s %s\n", op, data_buf);
            strcpy(input_buf_test_name, data_buf);
        } else if (strcmp(op, "RESULT") == 0 ) {
            printf("* Reading %s's Result\n", data_buf);
        } else if (strcmp(op, "USERID") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            memcpy(input_buf_uid, data_buf, USERID_BUFFER_SIZE);
        } else if (strcmp(op, "K") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            sm2_bigint_set_from_hex_str(&ctx->k, data_buf);
        } else if (strcmp(op, "PRIVATE_KEY") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            // Set Private Key
            sm2_ctx_set_private_key(ctx, data_buf);
        } else if (strcmp(op, "PUBLIC_KEY") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            // Parse public key
            const int len = (int) strlen(data_buf);
            const int offset = data_buf[0] == 0x04; // 0x04|X|Y| format
            if (len - offset != SM2_KEY_HEX_LEN * 2) {
                printf("ERROR: Public key length not correct\n");
                exit(1);
            }
            char key_buf[SM2_KEY_HEX_LEN * 2 + 1];
            key_buf[SM2_KEY_HEX_LEN * 2] = '\0';
            strncpy(key_buf, data_buf + offset, SM2_KEY_HEX_LEN * 2);
            // Set Public Key
            sm2_ctx_set_public_key(ctx, key_buf);
        } else if (strcmp(op, "DIGEST") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            memcpy(input_buf_digest, data_buf, 64);
            input_buf_digest[64] = '\0';
        } else if (strcmp(op, "R") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            memcpy(input_buf_r, data_buf, 64);
            input_buf_r[64] = '\0';
        } else if (strcmp(op, "S") == 0) {
            fgets(data_buf, buf_size, input);
            line_trim(data_buf);
            memcpy(input_buf_s, data_buf, 64);
            input_buf_s[64] = '\0';
        } else if (strcmp(op, "DER") == 0) {
            fgets(data_buf, buf_size, input);
            // ignored.
        } else if (strcmp(op, "MESSAGE") == 0) {
            strcpy(input_buf_msg_type, data_buf);
            int msg_max_len = buf_size;
            int msg_len = 0;
            for (;;) {
                char ch = fgetc(input);
                // if (msg_len == msg_max_len) { // realloc
                //     msg_max_len += buf_size;
                //     msg_buf = realloc(msg_buf, sizeof(char) * msg_max_len);
                //     if (msg_buf == NULL) {
                //         printf("Failed to allocate memory for message buffer\n");
                //         exit(1);
                //     }
                // }
                if (ch == '\n' || ch == '\r' || ch == EOF) {
                    input_buf_msg[msg_len] = '\0';
                    //  printf("Message size = %d\n%s\n", strlen(msg_buf), msg_buf);
                    break;
                }
                input_buf_msg[msg_len++] = ch;
            }
            break;
        }
    }
    fclose(input);
}


/**
 * Sign the data provided by the context.
 *
 * @param input_file the input file contains the data provided on the provider
 * @param output_file the columns needed to fill the Excel
 */
void run_sign_test(const char* input_file, const char* output_file) {
    // Init SM2 context
    sm2_ctx* ctx = sm2_ctx_shared();
    sm2_ctx_init(ctx);

    // Read data from files
    read_test_file(input_file, ctx);
    sm2_bitstring user_id = sm2_bitstring_wrap(sm2_storage_uid, 64, 0);
    sm2_bitstring_set_from_cstr(&user_id, input_buf_uid);
    // show_sm2_ctx(ctx);

    printf("* User Id hex\n");
    sm2_show_bitstring(&user_id);

    // Build Za String
    sm2_bitstring zstr = sm2_bitstring_wrap(sm2_storage_zstr, ZSTRING_BUFFER_SIZE, 0);
    sm2_build_zstr(ctx, &zstr, &user_id);
    // show_bitstring(&zstr);

    // SM3_H256
    sm3_hash(sm2_storage_za, zstr.data, zstr.bit_len);
    sm2_bitstring za = sm2_bitstring_wrap(sm2_storage_za, 32, 256);
    printf("* Za hex\n");
    sm2_show_bitstring(&za);

    int msg_byte_len = (int) strlen(input_buf_msg) / 2;
    printf("@ Message Digest (bytes = %d)\n", msg_byte_len);
    sm2_bitstring za_msg = sm2_bitstring_wrap(sm2_storage_message, MESSAGE_BUFFER_SIZE, 0);
    sm2_bitstring_set_from_bytes(&za_msg, za.data, za.byte_len);
    sm2_bitstring_set_from_hex_str(&za_msg, input_buf_msg);

    // Choose K
    // Fix a k, so the result can be reproduced
    if (sm2_bigint_is_zero(&ctx->k)) {
        sm2_bigint_set_from_hex_str(&ctx->k,
            "6CB28D99385C175C94F94E934817663FC176D925DD72B727260DBAAE1FB2F988"
        );
    }

    FILE* output = fopen(output_file, "w");
    if (output == NULL) {
        printf("Failed to open file to write: %s\n", input_file);
        return;
    }
    fprintf(output, "RESULT %s\n", input_buf_test_name);
    fprintf(output, "PUBLIC_KEY hex\n");
    uint8_t public_key[64];
    sm2_bigint_to_bytes(&ctx->Px, public_key);
    sm2_bigint_to_bytes(&ctx->Py, public_key + 32);
    sm2_write_bytes(output, public_key, 64);

    clock_t sign_measure_start = clock();
    // Calculate digest
    // bitstring* digest = sm2_digest(za, msg);
    sm3_hash(sm2_storage_ma, za_msg.data, za_msg.bit_len);
    sm2_bitstring digest = sm2_bitstring_wrap(sm2_storage_ma, 32, 256);
    printf("* Digest\n");
    sm2_show_bitstring(&digest);
    fprintf(output, "DIGEST hex\n");
    sm2_write_bitstring(output, &digest);

    // Sign digest with k
    if (sm2_bigint_is_zero(&ctx->k)) {
        ctx->sm2_change_k(ctx);
    }
    sm2_signature signature;
    int sign_res = 0;
    do {
        sign_res = sm2_sign_generate(ctx, &digest, &signature);
        if (0 == sign_res) {
            printf("* Sign Pass\n");
            sm2_show_signature(&signature);
            printf("* ASN.1 in DER =\n");
            sm2_signature_to_der((uint8_t*) &sm2_storage_der, &signature);
            sm2_show_bytes((uint8_t *)&sm2_storage_der, ASN1_DER_BYTE_LEN, false);
            fprintf(output, "DER hex\n");
            sm2_write_bytes(output, sm2_storage_der, ASN1_DER_BYTE_LEN);
            fprintf(output, "R hex\n");
            sm2_write_bytes(output, signature.r, 32);
            fprintf(output, "S hex\n");
            sm2_write_bytes(output, signature.s, 32);
        } else {
            // try another k
            ctx->sm2_change_k(ctx);
        }
    }  while(sign_res != 0);

    double benchmark = (double)(clock() - sign_measure_start) / CLOCKS_PER_SEC;
    printf("\n===================== Sign used %f sec ", benchmark);
    printf("========================\n\n");
    fclose(output);
}

/**
 * Verify the result generated.
 *
 * @param result_file the result file created by run_context_sign()
 */
void run_verify_test(const char* result_file) {
    // Init SM2 context
    sm2_ctx* ctx = sm2_ctx_shared();
    sm2_ctx_init(ctx);
    read_test_file(result_file, ctx);
    sm2_bitstring r = sm2_bitstring_wrap(sm2_storage_r, 32, 0);
    sm2_bitstring s = sm2_bitstring_wrap(sm2_storage_s, 32, 0);
    sm2_bitstring digest = sm2_bitstring_wrap(sm2_storage_ma, 32, 0);
    sm2_bitstring_set_from_hex_str(&r, input_buf_r);
    sm2_bitstring_set_from_hex_str(&s, input_buf_s);
    sm2_bitstring_set_from_hex_str(&digest, input_buf_digest);
    clock_t v_measure_start = clock();
    // show_sm2_ctx(ctx);
    printf("* Verification Starts \n");
    printf("  digest = ");
    sm2_show_bitstring(&digest);
    printf("  r = ");
    sm2_show_bitstring(&r);
    printf("  s = ");
    sm2_show_bitstring(&s);
    sm2_signature signature;
    for (int i = 0; i < 32; i++) {
        signature.r[i] = r.data[i];
        signature.s[i] = s.data[i];
    }
    int v_res = sm2_sign_verify(
        ctx,
        &digest,
        &signature
    );
    if (v_res == 0) {
        printf("[√] Verification passed\n");
    } else {
        printf("[x] Verification failed: %d\n", v_res);
    }
    double benchmark = (double)(clock() - v_measure_start) / CLOCKS_PER_SEC;
    printf("\n===================== Verify used %f sec ", benchmark);
    printf("=======================\n\n");
}

/*
// alternative params
void sm2_ctx_init_mod(sm2_ctx* ctx) {
    // Load sample params
    bigint_set_from_hex_str(ctx->p, "8542D69E4C044F18E8B92435BF6FF7DE457283915C45517D722EDB8B08F1DFC3");
    bigint_set_from_hex_str(ctx->a, "787968B4FA32C3FD2417842E73BBFEFF2F3C848B6831D7E0EC65228B3937E498");
    bigint_set_from_hex_str(ctx->b, "63E4C6D3B23B0C849CF84241484BFE48F61D59A5B16BA06E6E12D1DA27C5249A");
    bigint_set_from_hex_str(ctx->n, "8542D69E4C044F18E8B92435BF6FF7DD297720630485628D5AE74EE7C32E79B7");
    bigint_set_from_hex_str(ctx->Gx, "421DEBD61B62EAB6746434EBC3CC315E32220B3BADD50BDC4C4E6C147FEDD43D");
    bigint_set_from_hex_str(ctx->Gy, "0680512BCBB42C07D47349D2153B70C4E5D7FDFCBFA36EA1A85841B9E46E09A2");
}
*/

/**
 * Encrypt the data provided by the context.
 *
 * @param input_file the input file contains the data provided on the provider
 * @param output_file the columns needed to fill the Excel
 */
void run_encrypt_test(const char* input_file, const char* output_file) {
    // Init SM2 context
    sm2_ctx* ctx = sm2_ctx_shared();
    sm2_ctx_init(ctx);
    // sm2_ctx_init_mod(ctx);

    // Read data from files
    read_test_file(input_file, ctx);
    // sm2_show_ctx(ctx);

    // Done with Input
    sm2_bitstring plain_msg = sm2_bitstring_wrap(sm2_storage_message, MESSAGE_BUFFER_SIZE, 0);
    if (strcmp(input_buf_msg_type, "cstr") == 0) {
        sm2_bitstring_set_from_cstr(&plain_msg, input_buf_msg);
    } else {
        sm2_bitstring_set_from_hex_str(&plain_msg, input_buf_msg);
    }
    printf("* Plain Message hex = %u bits, %u bytes\n", plain_msg.bit_len, plain_msg.byte_len);
    sm2_show_bitstring(&plain_msg);

    FILE* output = fopen(output_file, "w");
    if (output == NULL) {
        printf("Failed to open file to write: %s\n", input_file);
        return;
    }
    fprintf(output, "RESULT %s\n", input_buf_test_name);
    fprintf(output, "PRIVATE_KEY hex\n");
    sm2_write_bigint(output, &ctx->d);
    clock_t encrypt_measure_start = clock();
    int encrypt_res = 0;
    sm2_bitstring encrypted_msg = sm2_bitstring_wrap(sm2_storage_encrypt, MESSAGE_BUFFER_SIZE, 0);
    do {
        encrypt_res = sm2_data_encrypt(ctx, &plain_msg, &encrypted_msg);
        if (0 == encrypt_res) {
            printf("@ Encrypted Data\n");
            sm2_show_bitstring(&encrypted_msg);
            fprintf(output, "MESSAGE hex\n");
            sm2_write_bitstring(output, &encrypted_msg);
        } else {
            ctx->sm2_change_k(ctx);
        }
    } while (encrypt_res != 0);

    double benchmark = (double)(clock() - encrypt_measure_start) / CLOCKS_PER_SEC;
    printf("\n===================== Encrypt used %f sec ", benchmark);
    printf("========================\n\n");
    fclose(output);
}

/**
 * Decrypt the data provided by the context.
 *
 * @param result_file the result file created by run_encrypt_test()
 */
void run_decrypt_test(const char* result_file) {
    // Init SM2 context
    sm2_ctx* ctx = sm2_ctx_shared();
    sm2_ctx_init(ctx);
    // sm2_ctx_init_mod(ctx);

    // Read data from files
    read_test_file(result_file, ctx);
    // show_sm2_ctx(ctx);

    // Done with Input
    sm2_bitstring encrypted_msg = sm2_bitstring_wrap(sm2_storage_message, MESSAGE_BUFFER_SIZE, 0);
    sm2_bitstring_set_from_hex_str(&encrypted_msg, input_buf_msg);
    printf("* Encrypted Message hex = %u bits, %u bytes\n", encrypted_msg.bit_len, encrypted_msg.byte_len);
    sm2_show_bitstring(&encrypted_msg);

    clock_t decrypt_measure_start = clock();
    sm2_bitstring plain_msg = sm2_bitstring_wrap(sm2_storage_encrypt, MESSAGE_BUFFER_SIZE, 0);
    int decrypt_res = sm2_data_decrypt(
        ctx,
        &encrypted_msg,
        &plain_msg
    );
    if (decrypt_res == 0) {
        printf("@ Decrypted Data\n");
        sm2_show_bitstring(&plain_msg);
        // plain_msg.data[plain_msg.byte_len] = '\0';
        // printf("@ Decrypted Message\n%s\n", plain_msg.data);
    } else {
        printf("[x] Data decryption failed: %d\n", decrypt_res);
    }

    double benchmark = (double)(clock() - decrypt_measure_start) / CLOCKS_PER_SEC;
    printf("\n===================== Decrypted used %f sec ", benchmark);
    printf("=======================\n\n");
}

int main(void) {
    // initializes the random number generator with the current time,
    // srand((unsigned int)time(NULL));
    // test_sm3_hash_1();
    // test_sm3_hash_2();
    // test_bigint_1();
    // test_sm2_ecc();
    // test_sm2_sign();
    // test_sm2_verify();

    run_sign_test(
        "../data/contest_case_1.txt",
        "../data/contest_case_1_out.txt"
    );
    run_verify_test(
        "../data/contest_case_1_out.txt"
    );
    run_sign_test(
        "../data/contest_case_2.txt",
        "../data/contest_case_2_out.txt"
    );
    run_verify_test(
        "../data/contest_case_2_out.txt"
    );
    run_sign_test(
        "../data/contest_case_3.txt",
        "../data/contest_case_3_out.txt"
    );
    run_verify_test(
        "../data/contest_case_3_out.txt"
    );
    // k will change
    run_sign_test(
        "../data/contest_case_4.txt",
        "../data/contest_case_4_out.txt"
    );
    run_verify_test(
        "../data/contest_case_4_out.txt"
    );

    // Encryption & Decryption
    run_encrypt_test(
        "../data/test_encrypt_1.txt",
        "../data/test_encrypt_1_out.txt"
    );
    run_decrypt_test(
        "../data/test_encrypt_1_out.txt"
    );
    printf("# All task finished\n");
    return 0;
}